import os

import cv2
import numpy as np

from tools.flip import compute_flip
from tools.utils import chw_to_hwc, index2color, get_magma_map, hwc_to_chw


def flip1(img1, img2, diff_image_path, monitor_distance=0.7, monitor_width=0.7, monitor_resolution_x=3840,
          use_color_map=True):
    # Compute number of pixels per degree of visual angle
    pixels_per_degree = monitor_distance * (monitor_resolution_x / monitor_width) * (np.pi / 180)
    # Compute FLIP map
    deltaE = compute_flip(img1, img2, pixels_per_degree)
    # Save error map
    index_map = np.floor(255.0 * deltaE.squeeze(0))
    if use_color_map:
        result = chw_to_hwc(index2color(index_map, get_magma_map()))
    else:
        result = index_map / 255.0
    # save_image(diff_image_path, result)
    result = (np.clip(result, 0.0, 1.0) * 255.0).astype(np.uint8)
    result = cv2.cvtColor(result, cv2.COLOR_BGR2GRAY)
    cv2.imwrite(diff_image_path, result)


def main():
    data_dir = 'E:\\workspace-pycharm\\deepfakes\\research-tools\\images'
    img1 = cv2.imread(os.path.join(data_dir, '1.png'), cv2.IMREAD_COLOR)
    img1 = np.asarray(img1).astype(np.float32)
    img1 = hwc_to_chw(img1) / 255.

    img2 = cv2.imread(os.path.join(data_dir, '2.png'), cv2.IMREAD_COLOR)
    img2 = np.asarray(img2).astype(np.float32)
    img2 = hwc_to_chw(img2) / 255.

    # img1 = load_image_array(os.path.join(data_dir, '1.png'))
    # img2 = load_image_array(os.path.join(data_dir, '2.png'))
    diff_image_path = os.path.join(data_dir, 'flip-opencv.png')
    flip1(img1, img2, diff_image_path)

    # diff = img1 - img2
    # diff = cv2.cvtColor(diff, cv2.COLOR_BGR2GRAY)
    # diff = diff / 255
    # diff[diff < 0.9] = 0
    # print(np.max(diff), np.min(diff))
    # diff = (diff * 255).astype(np.uint8)
    # flip(img1, img2, os.path.join(data_dir, 'diff3.png'))


if __name__ == '__main__':
    main()
